package gsar.client;


import searchandrescue.utils.Base64Coder;
import searchandrescue.utils.JavaUtils;
import gsar.client.GWMap.Mapcell;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.PrintWriter;
import java.io.StringReader;
import java.net.Socket;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import org.jdom.Document;
import org.jdom.Element;
import org.jdom.input.SAXBuilder;

/*
Copyright (c) 2009, Eric Eaton, University of Maryland Baltimore County
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of the University of Maryland Baltimore County nor the
      names of its contributors may be used to endorse or promote products
      derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS BE LIABLE FOR ANY
DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

/** The abstract class to represent a GSAR rescue agent.
 * 
 *  To create a rescue agent, create a subclass of the Client,
 *  implementing the abstract methods.
 *  
 *  Additionally, to make your agent runnable from the command line,
 *  write a main() method that calls Client.main().  Creating this 
 *  custom main() method also allows you to implement your own 
 *  command line argument handling, although the arguments to 
 *  Client.main() must be unchanged.  Here is an example main() 
 *  method for your reference:
 *  
 *  <code>
 *  public class MyClient extends Client {
 *  	...
 *  	public static void main(String[] args) {
 *  		main(new MyClient(), args);
 *  	}
 *  }
 *  </code>
 *  
 * @author Eric Eaton
 *
 */
public abstract class Client {	
	
	/** The agent's map of the gridworld */
	private GWMap m_map = null;
	
	/** the current simulation time, which starts at zero */
	private int m_currentTime = 0;

	/** Advances the simulation time one step
	 */
	private void step() { m_currentTime++; }
	
	/** Gets the current simulation time.
	 * @return the current simulation time
	 */
	protected int getCurrentTime() { return m_currentTime; }

	
	/** Set the agent's representation of the gridworld map.
	 * @param map the gridworld map
	 */
	protected void setMap(GWMap map) { m_map = map; }
	
	/** Get the agent's representation of the gridworld map.
	 * @return the gridworld map
	 */
	protected GWMap getMap() { return m_map; }
	
	
	/** This function is called automatically to allow the
	 *  agent to initialize itself.  The map has already
 	 *  been initialized when this function is called.
	 */
	protected abstract void initialize();
	
	
	/** This function allows the agent to process its final score.
	 *  @param score the agent's final score
	 */
	protected abstract void processResult(double score);

	
	/** Given a set of perceptions, it chooses the next action to take.
	 *  @param perceptions the current perceptions
	 *  @return an action structure defining the next action to take.
	 */
	protected abstract GWAction chooseAction(Perceptions perceptions);

	
	
	/** The helper main function for subclasses.  This function runs the
	 * rescue agent, interacts with the GSAR server, and completes a single
	 * episode of search and rescue.
	 * 
	 * @param agent the implemented GSAR rescue agent
	 * @param args the command line arguments, which are of the form:
	 *             "<server address> <server port number> <agent name> [<image filename>]"
	 */
	protected static void main(Client agent, String[] args) {	
		
		// parse the command line arguments
		String serverAddress = "";
		int serverPortNumber = 0;
		String name = "";
		String imageFilename = null;
		
		try {
			serverAddress = args[0];
			serverPortNumber = Integer.parseInt(args[1]);
			name = args[2];
			try {
				imageFilename = args[3];
			} catch (Exception e) {}
		} catch (Exception e) {
			System.out.println("Usage:  <server address> <server port number> <agent name> [<image filename>]");
			System.exit(1);
		}

		
		
		try {	
			// connect to the server
			try {
				m_serverConnection = new Socket(serverAddress, serverPortNumber);
			} catch (UnknownHostException e) {
				throw new IllegalArgumentException("Invalid server address.");
			} catch (IOException e) {
				throw new Exception("Could not connect to server "+serverAddress+":"+serverPortNumber);
			}
			
			// get the comm streams to and from the server
			m_serverInStream = new BufferedReader(new InputStreamReader(m_serverConnection.getInputStream()));
			m_serverOutStream = new PrintWriter(m_serverConnection.getOutputStream(), true);
			
			
			
			// register the client agent with the server
			String clientRegistrationXML = getClientRegistrationXML(name, imageFilename);
			m_serverOutStream.println(clientRegistrationXML);
			m_serverOutStream.flush();
			
			System.out.println("Connected to the simulation server at "+serverAddress+":"+serverPortNumber+".");
			
			
			
			try {
				
				// receive and process the initial knowledge
				GWMap map = processInitialKnowledge(m_serverInStream.readLine());
				agent.setMap(map);
					
				// let the agent initialize itself
				agent.initialize();

				//loop repeatedly in the perception-action cycle
				while (true) {
					
					// get the next perception until the perception time matches the simulation time
					String perception = null;
					int skipCount = -1;
					do {
						perception = m_serverInStream.readLine().trim();
						skipCount++;
					} while (m_serverInStream.ready());
					
					if (skipCount > 0) {
						System.out.println("Agent got behind "+skipCount+" timesteps, so it has discarded those perceptions.");
					}
						
					// check for and process the final score
					if (perception.startsWith("<finalscore>")) {
						agent.processResult(parseScoreFromXML(perception));
						break;
					}
					
					// parse the perception, get the next action, and return it
					GWAction currentAction = agent.chooseAction(parsePerceptionsFromXML(perception));
					m_serverOutStream.println(actionToXML(currentAction));
					m_serverOutStream.flush();
					
					// increment the current time
					agent.step();
				}
				
			} catch (IOException e) {
				System.out.println("Connection lost.");
			}
		} catch (UnknownHostException e) {
			System.out.println("Unknown host.");
		} catch (IOException e) {
			System.out.println("Could not connect to given host.");
		} catch (Exception e) {
			System.out.println(e.getMessage());
			e.printStackTrace();
		} finally {
			// shutdown the connections
			try {
				m_serverOutStream.close();
				m_serverInStream.close();
				m_serverConnection.close();
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
	}

	private static Socket m_serverConnection;

	/** The input stream to the server. */
	private static BufferedReader m_serverInStream;

	/** The output stream from the server. */
	private static PrintWriter m_serverOutStream;

	
	/** Converts an action to xml that can be understood by the server.
	 * @param action the action
	 * @return a string containing xml data that can be sent to the server
	 */
	private static String actionToXML (GWAction action) {
		String xmlStr = "<action> <name> "+action.getName()+" </name>";
		GWDirection actionDirection = action.getDirection();
		if (actionDirection != null) {
			xmlStr += " <direction> "+actionDirection+" </direction>";
		}
		xmlStr += " </action>";
		return xmlStr;
	}
	
	/** Construct the xml for registering an agent with the server.
	 * @param agentName the name of the rescue agent
	 * @param imageFilename the image representing that agent
	 * @return the xml registration string to send to the server
	 */
	private static String getClientRegistrationXML(String agentName, String imageFilename) {
		// encode the image if it is specified
		String encodedImage = null;
		if (imageFilename != null) {
			// encode the image
			try {
				byte[] imageFileAsBytes = JavaUtils.getBytesFromFile(new File(imageFilename));
				encodedImage = new String(Base64Coder.encode(imageFileAsBytes));
			} catch (IOException e) {
				System.err.println("Error:  could not read image \""+imageFilename+"\":  "+e.getMessage());
				encodedImage = null;
			}
		}
		
		// construct the registration string
		String agentRegistrationXML = "<register_agent> <name> "+agentName+" </name>";
		if (encodedImage != null) {
			agentRegistrationXML += " <image> "+encodedImage+" </image>";
		}
		agentRegistrationXML += " </register_agent>";
		
		return agentRegistrationXML;
	}
	
	/** Retrieves the root element, confirming that it is of the correct type.
	 * @param xmlInput the xml string
	 * @param rootElementName the name of the root element
	 * @return the root element
	 * @throws Exception if the xmlInput doesn't match the given root element name
	 */
	private static Element getRootElement(String xmlInput, String rootElementName) throws Exception {
		// convert the command string to the action
		SAXBuilder parser = new SAXBuilder();
	    Document cmd;
		try {
			cmd = parser.build(new StringReader(xmlInput));
		} catch (Exception e) {
			System.err.println("Cannot parse xml: "+xmlInput);
			throw e;
		}
		
		// confirm that the client sent the proper command
		if (!cmd.getRootElement().getName().equalsIgnoreCase(rootElementName)) {
			System.err.println("The xml root element is not a "+rootElementName+" element: "+xmlInput);
			throw new Exception("The xml root element is not a "+rootElementName+" element: "+xmlInput);
		}
	    
	    // get the root element
	    return cmd.getRootElement();
	}
	
	
	

	/**
	 * Parses the initial map xml data sent from the server.  This will
	 * include static information about the world (map layout, goal
	 * locations).  Data is expected in the following form:
	 *
	 * <map>
	 *   ...
	 * </map>
	 *
	 * @param xmlInput an xml string
	 * @return a map of the gridworld as represented by the initial knowledge
	 */
	private static GWMap processInitialKnowledge(String xmlInput) {
		Element node = null;
		try {
			node = getRootElement(xmlInput, "map");
		} catch (Exception e) {
			System.err.println(e.getMessage());
			return null;
		}
	
		// parse the XML map spec into a list of mapcells
		List<GWMap.Mapcell> mapcells = parseMapFromXML(node);
	
		// initialize the map from the parsed XML map spec
		GWMap map = new GWMap(mapcells);
	
		return map;
	}	
	
	
	








    /*
     *=========================================================================
	 *=========================================================================
	 *
	 * XML Parser Functions
	 *
	 *=========================================================================
	 *=========================================================================
	 */

    /** Parses the final score information sent by the server at the
	 *  end of the simulation.  Data is expected in the format:
	 *  <finalscore></finalscore>
     * @param input an xml string
	 * @return the final score
	 * @throws Exception if the score was not able to be parsed
	 */
	private static double parseScoreFromXML (String xmlInput) throws Exception {
		try {
			// get the score
			String scoreStr = getRootElement(xmlInput, "finalscore").getValue().trim();
			return Double.parseDouble(scoreStr);
		} catch (Exception e) {
			throw new Exception("Unable to parse the score: "+e.getMessage());
		}
	}


	/** Parses the data sent from the server at each time step of the
	 *  simulation.  This will include the status of the agent's last
	 *  action, and information about the area visible to the agent.
	 *  Data is expected in the following form:
	 *  <perceptions>
	 *    <action>
	 *       <name>NAME</name>
	 *       <direction>DIR</direction>
	 *       <time>TIME</time>
	 *       <status>STATUS</status>
	 *     </action>
	 *     <agent>
	 *       <coordinates><x>X</x><y>Y</y></coordinates>
	 *       <name>NAME</name>
	 *       <description>DESC</description>
	 *     </agent>
	 *     <cell>
	 *       <coordinates><x>X</x><y>Y</y></coordinates>
	 *       <marker>foo</marker>
	 *       <wall>SOUTH</wall>
	 *       ...
	 *       <occupant>
	 *         <object>
	 *           <coordinates><x>X</x><y>Y</y></coordinates>
	 *           <name>NAME</name>
	 *           <description>DESC</description>
	 *         </object>
	 *       </occupant>
	 *     </cell>
	 *     ...
	 *   </perceptions>
	 *  
	 * @param xmlInput an xml string
	 * @return a perceptions structure representing the data
	 */
	private static Perceptions parsePerceptionsFromXML (String xmlInput) {
		Element rootNode = null;
		try {
			rootNode = getRootElement(xmlInput, "perceptions");
		} catch (Exception e) {
			System.err.println(e.getMessage());
			return null;
		}
		
		GWActionStatus actionStatus = parseActionStatusFromXML(rootNode.getChild("action"));
		GWAgent self = parseAgentFromXML(rootNode.getChild("agent"));
		
		ArrayList<GWCell> longRangeSensors = new ArrayList<GWCell>();
		List<Element> cellElements = rootNode.getChildren("cell");
		for (Element cellElement : cellElements) {
			longRangeSensors.add(parseCellFromXML(cellElement));
		}
		
		GWObject shortRangeSensors = null;
		if (actionStatus != null && actionStatus.getStatus() instanceof GWObject) {
			shortRangeSensors = (GWObject) actionStatus.getStatus();
		}
		
		return new Perceptions(actionStatus, self, shortRangeSensors, longRangeSensors);
	}
		 

	




	/** Creates a list of goals from an xml goal specification in the S-XML DOM
	 * of the format 
	 *    <goals>
	 *      <coordinates> ... </coordinates>
	 *      <coordinates> ... </coordinates>
	 *      ...
	 *    </goals>
	 *
 	 * @param node the parent node in the DOM
	 * @return a list of the goals contained in the node
	 */
	private static GWCoordinate[] parseGoalsFromXML(Element node) {
		if (!node.getName().equalsIgnoreCase("goals"))
			return null;
			
		// the goal information will be processed into a list of coordinates
		// structures, one for each goal location
		List<Element> coordinatesElements = node.getChildren("coordinates");
		GWCoordinate[] goals = new GWCoordinate[coordinatesElements.size()];
		
		int i=0;
		for (Element e : coordinatesElements) {
			goals[i++] = parseCoordinatesFromXML(e);
		}
		
		return goals;
	}
	
	/** Parses a cell from an S-XML DOM.
	 * Data is expected in the following form:
 	 *
	 *   <cell>
	 *     <coordinates> ... </coordinates>
	 *     <marker> m </marker>
	 *     <wall>NORTH</wall>
	 *     <wall>EAST</wall>
	 *     <occupant>...</occupant>
	 *   </cell>
	 *
  	 * @param node the cell node
	 * @return the cell
	 */
	private static GWCell parseCellFromXML(Element node) {
		if (!node.getName().equalsIgnoreCase("cell"))
			return null;
		
		List<Element> wallElements = node.getChildren("wall");
		Set<String> walls = new HashSet<String>();
		for (Element e : wallElements) {
			walls.add(e.getValue());
		}
		
		String marker = null;
		Element markerNode = node.getChild("marker");
		if (markerNode != null) {
			marker = markerNode.getValue();
		}
		
		return new GWCell(parseCoordinatesFromXML(node.getChild("coordinates")),
						  walls,
						  parseOccupantFromXML(node.getChild("occupant")),
						  marker);
	}
	

	/** Retrieves the structure representing the contents of the occupant node in the S-XML DOM.
	 * @param node the occupant node in the DOM
	 * @return the structure representing the occupant
	 */
	private static GWObject parseOccupantFromXML(Element node) {
		if (node == null || !node.getName().equalsIgnoreCase("occupant"))
			return null;
		
		Element firstChild = (Element) node.getChildren().get(0);
		
		GWObject obj = parseObjectFromXML(firstChild);
		if (obj != null)
			return obj;
		
		return parseAgentFromXML(firstChild);
	}


	/** Creates an action-status structure from the first action-status node under a given parent in the S-XML DOM.
	 *
	 * @param node the parent node in the DOM
	 * @return a new action-status structure populated from the DOM node, or NULL if no matching node was found
	 */
	private static GWActionStatus parseActionStatusFromXML(Element node) {
		if (node == null || !node.getName().equalsIgnoreCase("action"))
			return null;
		
		// check if the status field has a more complex structure inside it
		Element statusNode = node.getChild("status");
		
		Object status = null;
		// if the statusNode doesn't have any children, its value is a primitive
		if (statusNode.getChildren().size() == 0) {
			status = statusNode.getValue();
		} else {
			Element statusChild = (Element) statusNode.getChildren().get(0);
			
			// try to parse out an object
			status = parseObjectFromXML(statusChild);
			
			// if that didn't work, try to parse out an agent
			if (status == null) {
				status = parseAgentFromXML(statusChild);
			}
			
			// if that didn't work, the status structure is unknown
			status = statusNode.getValue();
			
		}
		
		Element directionNode = node.getChild("direction");
		String directionStr = null;
		if (directionNode != null) {
			directionStr = directionNode.getValue();
		}
		
		GWAction action = new GWAction(node.getChild("name").getValue(),
				                       directionStr);
		
		return new GWActionStatus(action,
				                Integer.parseInt(node.getChild("time").getValue()),
				                status);
	}



	/** Creates an agent structure from the first agent node under a given parent in the S-XML DOM.
	 * @param node the parent node in the DOM
   	 * @return a new agent structure populated from the DOM node, or NULL if no matching node was found
 	 */
	private static GWAgent parseAgentFromXML(Element node) {
		  if (node == null || !node.getName().equalsIgnoreCase("agent"))
			  return null;
		  
		  Element descriptionNode = node.getChild("description");
		  Object[] description = null;
		  if (descriptionNode != null) {
			  description = parseFeatureVectorFromXML(descriptionNode.getChild("featurevector"));
		  }
		  
		  return new GWAgent(node.getChild("name").getValue(),
				             parseCoordinatesFromXML(node.getChild("coordinates")),
					             description);
	  }


	/** Creates an object structure from the first object node under a given parent in the S-XML DOM.
	 *
	 * @param node the parent node in the DOM
	 * @returns a new object structure populated from the DOM node, or NULL if no matching node was found
	 */
	private static GWObject parseObjectFromXML(Element node) {
		if (node == null || !node.getName().equalsIgnoreCase("object"))
			  return null;
		  
		Element descriptionNode = node.getChild("description");
		Object[] description = null;
		if (descriptionNode != null) {
			description = parseFeatureVectorFromXML(descriptionNode.getChild("featurevector"));
		}
	  
		return new GWObject(node.getChild("name").getValue(),
				            parseCoordinatesFromXML(node.getChild("coordinates")),
			                description);
	}


	/* Parses a feature vector from an S-XML DOM.
	 * Data is expected in the following form:
	 *
	 *   <featurevector>
	 *     <value> value1 </value>
	 *     <value> value2 </value>
	 *     ...
	 *   </featurevector>
	 *
	 * @param node the cell node
	 * @returns a list of the feature vector's values in-order
	 */
    private static Object[] parseFeatureVectorFromXML(Element node) {
	    if (!node.getName().equalsIgnoreCase("featurevector")) {
	    	return null;
	    }
	    
	    List<Element> valueNodes = node.getChildren("value"); 
	    Object[] values = new Object[valueNodes.size()]; 
	    int i=0;
	    for (Element e : valueNodes) {
	    	String eValueStr = e.getValue();
	    	// try making it an Integer
	    	try {
	    		values[i] = Integer.parseInt(eValueStr);
	    	} catch (Exception e1) {
	    		// try making it a Double
	    		try {
	    			values[i] = Double.parseDouble(eValueStr);
	    		} catch (Exception e2) {
	    			// store it as a string
	    			values[i] = eValueStr;
	    		}
	    	}
	    	
	    	
	    	i++;
	    }
	    return values;
	}



	/**
  	 * Creates a coordinates structure from a coordinates element in the S-XML DOM.
	 * @param node the coordinates node from the DOM
	 * @return a new cell structure populated with the values from the DOM node
	 */	
	private static GWCoordinate parseCoordinatesFromXML(Element node) {
		 if (!node.getName().equalsIgnoreCase("coordinates"))
			 return null;
		 
		 return new GWCoordinate(Integer.parseInt(node.getChild("x").getValue()),
				 		  		  Integer.parseInt(node.getChild("y").getValue()));
	}
	
	
	/** Parses the map from an S-XML DOM.
	 * Data is expected in the following form:
 	 *
     *   <map>
	 *     <mapcell> ... </mapcell>
	 *     <mapcell> ... </mapcell>
	 *     ...
	 *   </map>
	 * 
	 * @param node the map node
	 * @return a list of mapcells
	 */
	private static List<Mapcell> parseMapFromXML (Element node) {
		if (!node.getName().equalsIgnoreCase("map"))
			return null;
		
		List<Mapcell> mapcells = new ArrayList<Mapcell>();
		List<Element> mapcellElements = node.getChildren("mapcell");
		for (Element e : mapcellElements) {
			mapcells.add(parseMapcellFromXML(e));
		}
		return mapcells;
	}

	/** 
	 * Parses a mapcell from an S-XML DOM.
	 * Data is expected in the following form:
	 *
	 *   <mapcell id="cell_ID_0">
	 *     <coordinates> ... </coordinates>
	 *     <marker> m </marker>
	 *     <adjacentcells>
	 *       <mapcell id="cell_ID_1"/>
	 *       <mapcell id="cell_ID_2"/>
	 *     </adjacentcells>
	 *   </mapcell>
	 *
	 * @param node the cell node
	 * @return a mapcell structure of this cell
	 */
	private static Mapcell parseMapcellFromXML(Element node) {
		if (!node.getName().equalsIgnoreCase("mapcell"))
			return null;
		
		// get a list of all the cells in the adjacent cells list
		List<Element> adjacentCells = node.getChild("adjacentcells").getChildren("mapcell");
		
		// get the ids of adjacent cells
		List<String> adjacentCellIDs = new ArrayList<String>();
		for (Element e : adjacentCells) {
			adjacentCellIDs.add(e.getAttribute("id").getValue());
		}
		
		String marker = null;
		Element markerNode = node.getChild("marker");
		if (markerNode != null) {
			marker = markerNode.getValue();
		}
		
		Mapcell mapcell = new Mapcell(parseCoordinatesFromXML(node.getChild("coordinates")),
									  node.getAttribute("id").getValue(),
									  adjacentCellIDs,
									  marker);
		return mapcell;
	}
}